use aes::Aes256;
use chacha20poly1305::ChaCha20Poly1305;
use chacha20poly1305::aead::{Aead, KeyInit, Payload};
use chacha20poly1305::Key;
use cipher::generic_array::GenericArray;
use cipher::BlockDecrypt;

pub fn decrypt_aes(file_data: &[u8], key: &[u8; 32]) -> Vec<u8> {
    let cipher = Aes256::new(GenericArray::from_slice(key));
    let mut decrypted_data = Vec::new();

    for chunk in file_data.chunks(16) {
        let mut block = [0u8; 16];
        block[..chunk.len()].copy_from_slice(chunk);
        let mut decrypted_block = block;
        cipher.decrypt_block(GenericArray::from_mut_slice(&mut decrypted_block));
        decrypted_data.extend_from_slice(&decrypted_block);
    }

    // Remove padding (PKCS7)
    if let Some(&padding_size) = decrypted_data.last() {
        let len = decrypted_data.len();
        if padding_size as usize <= len {
            decrypted_data.truncate(len - padding_size as usize);
        }
    }

    decrypted_data
}

pub fn decrypt_chacha(file_data: &[u8], key: &[u8; 32], nonce: &[u8; 12]) -> Vec<u8> {
    let cipher = ChaCha20Poly1305::new(Key::from_slice(key));
    let payload = Payload {
        msg: file_data,
        aad: b"optional_data",
    };

    match cipher.decrypt(nonce.into(), payload) {
        Ok(decrypted_data) => decrypted_data,
        Err(_) => {
            eprintln!("ChaCha decryption failed");
            Vec::new()
        }
    }
}